import pymysql
import pandas as pd


class UserMapper:
    # 初始化连接
    def __init__(self):
        self.host = "localhost"
        self.user = "root"
        self.password = "123456"
        self.database = "student_portrait_system"
        self.connection = pymysql.connect(
            host=self.host,
            user=self.user,
            password=self.password,
            database=self.database
        )

    # 读取所有用户数据，返回dataframe格式
    def read_all_users(self) -> pd.DataFrame:
        with self.connection.cursor() as cursor:
            sql = "SELECT * FROM users"
            cursor.execute(sql)
            result = cursor.fetchall()
            columns = [i[0] for i in cursor.description]
            df = pd.DataFrame(result, columns=columns)
            return df

    # 读取指定用户数据，返回dataframe格式
    def read_user(self, user_id):
        with self.connection.cursor() as cursor:
            sql = "SELECT * FROM users WHERE user_id = %s"
            cursor.execute(sql, (user_id,))
            result = cursor.fetchone()
            columns = [i[0] for i in cursor.description]
            df = pd.DataFrame([result], columns=columns)
            return df

    # 新增
    def add_user(self, username, password):
        with self.connection.cursor() as cursor:
            sql = "INSERT INTO users (username,password) VALUES (%s,%s)"
            cursor.execute(sql, (username, password))
            self.connection.commit()
            return True

    # 更新用户信息
    def update_user(self, user_id, username, password):
        with self.connection.cursor() as cursor:
            sql = "UPDATE users SET username = %s, password = %s WHERE user_id = %s"
            cursor.execute(sql, (username, password, user_id))
            self.connection.commit()
            return True

    # 删除用户
    def delete_user(self, user_id):
        with self.connection.cursor() as cursor:
            sql = "DELETE FROM users WHERE user_id = %s"
            cursor.execute(sql, (user_id,))
            self.connection.commit()
            return True

    # 登录校验
    def login(self, username, password):
        with self.connection.cursor() as cursor:
            sql = "SELECT * FROM users WHERE username = %s AND password = %s"
            cursor.execute(sql, (username, password))
            result = cursor.fetchone()
            if result:
                return True
            else:
                return False

    # 注册用户
    def register(self, username, password):
        # 先判断用户是否存在
        with self.connection.cursor() as cursor:
            sql = "SELECT * FROM users WHERE username = %s"
            cursor.execute(sql, (username,))
            result = cursor.fetchone()
            # 用户不存在
            if result:
                return False
            else:
                # 新增用户
                self.add_user(username, password)
                # 注册成功
                return True


    # 关闭连接
    def close_connection(self):
        self.connection.close()
